import glob
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from torchvision import datasets
import config as c
# 该库用于实现Python的自然排序
from natsort import natsorted
import matplotlib.pyplot as plt



def to_rgb(image):
    '''
    如果原始图像已经是RGB模式,paste方法不会改变图像的颜色模式。
    如果原始图像是其他模式,如RGBA或灰度, paste方法会自动将其转换为RGB模式
    '''
    rgb_image = Image.new("RGB", image.size)
    rgb_image.paste(image)
    return rgb_image


# 创建数据集
class Hinet_Dataset(Dataset):
    def __init__(self, transforms=None, mode="train"):

        self.transform = transforms
        self.mode = mode
        if mode == 'train':
            # train
            # glob函数用于查找目标路径下的所有png图片，返回一个列表，保存文件路径
            # 接着进行自然排序
            self.files = natsorted(sorted(glob.glob(c.TRAIN_PATH + "/*." + c.format_train)))
        else:
            # valid验证集
            self.files = sorted(glob.glob(c.VAL_PATH + "/*." + c.format_val))

        # 计算文件数量
        self.file_len = len(self.files)
        print("当前模式为：{}，读取了{}张图片".format(self.mode, self.file_len))

    def __getitem__(self, index):
        try:
            image = Image.open(self.files[index])
            image = to_rgb(image)
            item = self.transform(image) # 图像增广
            return item

        except:
            return self.__getitem__(index + 1)

    def __len__(self):
        if self.mode == 'shuffle':  # ？？？
            return max(len(self.files_cover), len(self.files_secret))

        else:
            return self.file_len


# 图像增广
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(), 
    transforms.RandomVerticalFlip(),
    transforms.RandomCrop(c.cropsize), # 随机裁剪
    transforms.ToTensor()
])

transform_val = transforms.Compose([
    transforms.CenterCrop(c.cropsize_val),  # 中心裁剪
    transforms.ToTensor(),
])




# 用函数对加载数据进行封装，避免全部加载

# Training data loader
def get_trainloader():
    trainloader = DataLoader(
        Hinet_Dataset(transforms=transform_train, mode="train"),
        batch_size=c.batch_size,
        shuffle=True,
        pin_memory=True,  # 将数据加载到cuda中
        # num_workers=8,
        drop_last=True
    )
    return trainloader


# Test data loader
def get_testloader():
    testloader = DataLoader(
        Hinet_Dataset(transforms=transform_val, mode="val"),
        batch_size=c.batchsize_val,
        shuffle=False,
        pin_memory=True,
        # num_workers=1,
        drop_last=True
    )
    return testloader



# 展示数据, 看看数据集能不能正常读取图片(测试完，是ok的)
# def im_convert(tensor):
#     image = tensor.to("cpu").clone().detach()
#     image = image.numpy().squeeze()
#     image = image.transpose(1,2,0)
#     image = image.clip(0,1)
    
#     return image

# fig=plt.figure(figsize=(20, 12))
# columns = 1
# rows = 2
# batch_img = next(iter(testloader))
# print(batch_img.shape)

# for idx in range(columns*rows):
#     ax = fig.add_subplot(rows, columns, idx+1)
#     plt.imshow(im_convert(batch_img[idx]))
# plt.show()

